package com.gzist.springailab06.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.util.MimeType;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.nio.file.*;
import java.util.Comparator;
import java.util.List;

/**
 * 多模态能力实践接口
 */
@RestController
@RequestMapping("/api/chat")
public class MultiModalChatController {

    private final ChatClient chatClient;

    public MultiModalChatController(ChatModel chatModel) {
        this.chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(
                        new MessageChatMemoryAdvisor(new InMemoryChatMemory())
                ).build();
    }

    /**
     * 处理多模态聊天请求的接口方法
     *
     * @param text 文本输入内容，必填参数
     * @param file 上传的文件附件唯一，可选参数
     * 支持的文件类型：图片、音频、文档，由于对接模型厂商的模型限制，目前仅支持图片和内容较少的文档。
     */
    @PostMapping("/multi-modal")
    public ResponseEntity<String> multiModalChat(
            @RequestParam String text,
            @RequestParam(required = false) MultipartFile file) {
        try {
            // 处理多模态请求并获取响应
            String response = processMultiModalRequest(text, file);
            return ResponseEntity.ok(response);
        } catch (IllegalArgumentException e) {
            // 处理非法参数异常
            return ResponseEntity.badRequest().body(e.getMessage());
        } catch (FileSizeLimitExceededException e) {
            // 处理文件大小超限异常
            return ResponseEntity.status(HttpStatus.PAYLOAD_TOO_LARGE)
                    .body("文件大小超过限制: " + e.getMessage());
        } catch (Exception e) {
            // 处理其他未知异常
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body("处理请求时发生错误: " + e.getMessage());
        }
    }

    private String processMultiModalRequest(String text, MultipartFile file) throws Exception {
        if (file == null || file.isEmpty()) {
            return chatClient.prompt().user(text).call().content();
        }
        String fileExtension = getFileExtension(file.getOriginalFilename());
        MimeType mimeType = MimeTypeUtils.parseMimeType(getMimeType(fileExtension));
        if (isImage(fileExtension) || isAudio(fileExtension)) {
            return processMediaContent(text, file, mimeType);
        } else if (isDocument(fileExtension)) {
            return processDocumentContent(text, file);
        } else {
            throw new IllegalArgumentException("不支持的文件类型: " + fileExtension);
        }
    }

    private String processMediaContent(String text, MultipartFile file, MimeType mimeType) throws IOException {
        // 确保文件资源可访问
        File tempFile = File.createTempFile("upload-", ".tmp");
        file.transferTo(tempFile);
        try {
            Resource resource = new FileSystemResource(tempFile);
            return chatClient.prompt()
                    .user(u -> u.text(text).media(mimeType, resource))
                    .call()
                    .content();
        } finally {
            Files.deleteIfExists(tempFile.toPath()); // 清理临时文件
        }
    }

    private String processDocumentContent(String text, MultipartFile file) throws Exception {
        Path tempDir = Files.createTempDirectory("spring-ai-docs");
        try {
            String fileName = file.getOriginalFilename();
            String uniqueFileName = "uploaded-" + System.currentTimeMillis() +
                    "." + getFileExtension(fileName);
            Path tempPath = tempDir.resolve(uniqueFileName);
            // 使用transferTo确保文件正确写入
            file.transferTo(tempPath);
            // 使用更可靠的文件资源处理
            Resource resource = new FileSystemResource(tempPath.toAbsolutePath().toString());
            // 读取文档内容
            String docText = readDocumentContent(resource);
            // 调用大模型处理
            return chatClient.prompt()
                    .user(text + "\n\n" + docText)
                    .call()
                    .content();
        } finally {
            deleteDirectory(tempDir);
        }
    }

    private String readDocumentContent(Resource resource) throws IOException {
        // 改进的文档读取逻辑（需确保支持所需格式）
        List<Document> documents = new TextReader(resource).read();
        StringBuilder content = new StringBuilder("文档内容：\n");
        for (Document doc : documents) {
            content.append(doc.getFormattedContent()).append("\n\n");
        }
        return content.toString();
    }

    /**
     * 删除临时目录及其内容
     */
    private void deleteDirectory(Path directory) {
        try {
            if (Files.exists(directory)) {
                Files.walk(directory)
                        .sorted(Comparator.reverseOrder())
                        .forEach(path -> {
                            try {
                                Files.deleteIfExists(path);
                            } catch (IOException e) {
                                System.err.println("无法删除文件: " + path + ", 错误: " + e.getMessage());
                            }
                        });
            }
        } catch (IOException e) {
            System.err.println("删除临时目录失败: " + directory + ", 错误: " + e.getMessage());
        }
    }

    private String getFileExtension(String fileName) {
        if (fileName == null) return "";
        int dotIndex = fileName.lastIndexOf('.');
        return dotIndex > 0 ? fileName.substring(dotIndex + 1).toLowerCase() : "";
    }

    private String getMimeType(String extension) {
        return switch (extension) {
            case "png" -> "image/png";
            case "jpg", "jpeg" -> "image/jpeg";
            case "gif" -> "image/gif";
            case "mp3" -> "audio/mp3";
            case "wav" -> "audio/wav";
            case "pdf" -> "application/pdf";
            case "doc" -> "application/msword";
            case "docx" -> "application/vnd.openxmlformats-officedocument.wordprocessingml.document";
            case "txt" -> "text/plain";
            default -> "application/octet-stream";
        };
    }

    private boolean isImage(String extension) {
        return List.of("png", "jpg", "jpeg", "gif").contains(extension);
    }

    private boolean isAudio(String extension) {
        return List.of("mp3", "wav").contains(extension);
    }

    private boolean isDocument(String extension) {
        return List.of("pdf", "doc", "docx", "txt").contains(extension);
    }

    // 自定义异常类
    public static class FileSizeLimitExceededException extends Exception {
        public FileSizeLimitExceededException(String message) {
            super(message);
        }
    }
}